﻿#include "http_struct.h"

namespace beast = boost::beast;
namespace http = beast::http;
namespace net = boost::asio;
using tcp = boost::asio::ip::tcp;


// 异步处理器包装器
template<typename Handler>
class AsyncHandlerWrapper {
private:
    Handler handler_;

public:
    explicit AsyncHandlerWrapper(Handler h) : handler_(std::move(h)) {
    }

    std::future<HttpResponse> operator()(const RequestContext &ctx) {
        std::promise<HttpResponse> promise;
        auto future = promise.get_future();

        // 在线程池中执行
        boost::asio::post(getThreadPool(), [this, ctx, promise = std::move(promise)]() mutable {
            try {
                HttpResponse response = handler_(ctx);
                promise.set_value(std::move(response));
            } catch (...) {
                promise.set_exception(std::current_exception());
            }
        });

        return future;
    }

private:
    static boost::asio::thread_pool &getThreadPool() {
        static boost::asio::thread_pool pool(std::thread::hardware_concurrency());
        return pool;
    }
};

// 改进的路由器
class AsyncRouter {
private:
    std::map<RouteKey, std::function<std::future<HttpResponse>(const RequestContext &)>, RouteKeyCompare> routes_;

public:
    // 添加异步处理器
    template<typename Handler>
    void addAsyncRoute(const std::string &path, HttpMethod method, ContentType contentType, Handler handler) {
        RouteKey key(path, method, contentType);
        routes_[key] = AsyncHandlerWrapper<Handler>(std::move(handler));
    }

    // 添加同步处理器（自动包装为异步）
    template<typename Handler>
    void addRoute(const std::string &path, HttpMethod method, ContentType contentType, Handler handler) {
        addAsyncRoute(path, method, contentType,
                      [handler](const RequestContext &ctx) -> HttpResponse {
                          return handler(ctx);
                      }
        );
    }

    // 重载运算符方便添加路由
    template<typename Handler>
    AsyncRouter &operator()(const std::string &path, HttpMethod method, ContentType contentType, Handler handler) {
        addRoute(path, method, contentType, handler);
        return *this;
    }

    template<typename Handler>
    AsyncRouter &operator()(const std::string &path, HttpMethod method, Handler handler) {
        addRoute(path, method, ContentType::ANY, handler);
        return *this;
    }

    std::future<HttpResponse> matchAndHandleAsync(const RequestContext &context) {
        // 精确匹配
        RouteKey exactKey(context.path, context.method, context.contentType);
        auto it = routes_.find(exactKey);
        if (it != routes_.end()) {
            return it->second(context);
        }

        // 通配符匹配
        RouteKey anyTypeKey(context.path, context.method, ContentType::ANY);
        it = routes_.find(anyTypeKey);
        if (it != routes_.end()) {
            return it->second(context);
        }

        // 返回 404
        std::promise<HttpResponse> promise;
        promise.set_value(HttpResponse(404, "Not Found"));
        return promise.get_future();
    }
};

// HTTP 会话类
class AsyncHttpSession : public std::enable_shared_from_this<AsyncHttpSession> {
private:
    tcp::socket socket_;
    beast::flat_buffer buffer_;
    http::request<http::string_body> req_;
    AsyncRouter &router_;

public:
    AsyncHttpSession(tcp::socket socket, AsyncRouter &router)
        : socket_(std::move(socket)), router_(router) {
    }

    void start() {
        readRequest();
    }

private:
    static boost::asio::thread_pool &getThreadPool() {
        static boost::asio::thread_pool pool(std::thread::hardware_concurrency());
        return pool;
    }

    void readRequest() {
        auto self = shared_from_this();
        http::async_read(socket_, buffer_, req_,
                         [self](beast::error_code ec, std::size_t bytes_transferred) {
                             boost::ignore_unused(bytes_transferred);
                             if (!ec) {
                                 self->handleRequest();
                             }
                         });
    }

    void handleRequest() {
        // 创建请求上下文（在 IO 线程中快速完成）
        RequestContext context = createContext();
        auto self = shared_from_this();

        // 异步处理业务逻辑
        auto future_response = router_.matchAndHandleAsync(context);

        // 使用 shared_future 可以避免移动问题
        auto shared_future = future_response.share();

        // 使用线程池处理
        boost::asio::post(getThreadPool(), [self, shared_future = std::move(shared_future)]() {
            try {
                // shared_future 可以被拷贝
                HttpResponse response = shared_future.get();

                // 回到 IO 线程发送响应
                boost::asio::post(self->socket_.get_executor(),
                                  [self, response = std::move(response)]() {
                                      self->sendResponse(std::move(response));
                                  });
            } catch (const std::exception &e) {
                boost::asio::post(self->socket_.get_executor(),
                                  [self, error_msg = std::string(e.what())]() {
                                      self->sendErrorResponse("Internal Server Error: " + error_msg);
                                  });
            } catch (...) {
                boost::asio::post(self->socket_.get_executor(),
                                  [self]() {
                                      self->sendErrorResponse("Unknown Internal Server Error");
                                  });
            }
        });
    }


    RequestContext createContext() {
        RequestContext context;
        context.path = std::string(req_.target());

        // 分离路径和查询参数
        size_t queryPos = context.path.find('?');
        std::string queryString;
        if (queryPos != std::string::npos) {
            queryString = context.path.substr(queryPos + 1);
            context.path = context.path.substr(0, queryPos);
        }

        context.queryParams = parseQueryParams(queryString);
        context.body = req_.body();

        // 解析 HTTP 方法
        if (req_.method() == http::verb::get) context.method = HttpMethod::GET;
        else if (req_.method() == http::verb::post) context.method = HttpMethod::POST;
        else if (req_.method() == http::verb::put) context.method = HttpMethod::PUT;
        else context.method = HttpMethod::GET;

        // 解析 Content-Type
        auto contentTypeIt = req_.find(http::field::content_type);
        std::string contentTypeStr = (contentTypeIt != req_.end()) ? std::string(contentTypeIt->value()) : "";
        context.contentType = parseContentType(contentTypeStr);

        // 复制头部信息
        for (auto &field: req_) {
            context.headers[std::string(field.name_string())] = std::string(field.value());
        }

        return context;
    }

    ContentType parseContentType(const std::string &contentTypeStr) {
        if (contentTypeStr.empty()) return ContentType::ANY;
        if (contentTypeStr.find("application/json") != std::string::npos)
            return ContentType::JSON;
        if (contentTypeStr.find("application/x-www-form-urlencoded") != std::string::npos)
            return ContentType::FORM;
        if (contentTypeStr.find("text/html") != std::string::npos)
            return ContentType::HTML;
        if (contentTypeStr.find("text/") != std::string::npos)
            return ContentType::TEXT;
        return ContentType::ANY;
    }

    std::map<std::string, std::string> parseQueryParams(const std::string &query) {
        std::map<std::string, std::string> params;
        if (query.empty()) return params;

        size_t pos = 0;
        while (pos < query.length()) {
            size_t ampPos = query.find('&', pos);
            if (ampPos == std::string::npos) ampPos = query.length();

            std::string param = query.substr(pos, ampPos - pos);
            size_t eqPos = param.find('=');
            if (eqPos != std::string::npos) {
                std::string key = param.substr(0, eqPos);
                std::string value = param.substr(eqPos + 1);
                params[key] = value;
            } else {
                params[param] = "";
            }

            pos = ampPos + 1;
        }
        return params;
    }

    void sendResponse(const HttpResponse &response) {
        http::response<http::string_body> res{static_cast<http::status>(response.statusCode), 11};
        res.set(http::field::server, "Async HTTP Server");
        res.set(http::field::content_type, "application/json");

        // 设置自定义头部
        for (const auto &header: response.headers) {
            res.set(header.first, header.second);
        }

        res.body() = response.body;
        res.prepare_payload();

        auto self = shared_from_this();
        auto sp = std::make_shared<http::response<http::string_body> >(std::move(res));
        http::async_write(socket_, *sp,
                          [self, sp](beast::error_code ec, std::size_t bytes_transferred) {
                              boost::ignore_unused(bytes_transferred);
                              if (!ec) {
                                  beast::error_code ec2;
                                  self->socket_.shutdown(tcp::socket::shutdown_send, ec2);
                              }
                          });
    }

    void sendErrorResponse(const std::string &message) {
        HttpResponse response(500, message);
        sendResponse(response);
    }
};

// 异步 HTTP 服务器类
class AsyncHttpServer {
private:
    AsyncRouter router_;
    net::io_context &ioc_;
    tcp::acceptor acceptor_;
    std::vector<std::thread> io_threads_;

public:
    AsyncHttpServer(net::io_context &ioc, tcp::endpoint endpoint)
        : ioc_(ioc), acceptor_(ioc) {
        beast::error_code ec;
        acceptor_.open(endpoint.protocol(), ec);
        if (ec) throw std::runtime_error("Failed to open acceptor");

        acceptor_.set_option(net::socket_base::reuse_address(true), ec);
        if (ec) throw std::runtime_error("Failed to set reuse address");

        acceptor_.bind(endpoint, ec);
        if (ec) throw std::runtime_error("Failed to bind");

        acceptor_.listen(net::socket_base::max_listen_connections, ec);
        if (ec) throw std::runtime_error("Failed to listen");
    }

    AsyncRouter &getRouter() { return router_; }

    void run() {
        std::cout << "Async Server listening on " << acceptor_.local_endpoint() << std::endl;

        // 启动多个 IO 线程
        unsigned thread_count = std::thread::hardware_concurrency();
        io_threads_.reserve(thread_count - 1); // 主线程也算一个

        for (unsigned i = 0; i < thread_count - 1; ++i) {
            io_threads_.emplace_back([this]() {
                ioc_.run();
            });
        }

        doAccept();
        ioc_.run(); // 主线程也运行 IO

        // 等待其他线程结束
        for (auto &thread: io_threads_) {
            if (thread.joinable()) {
                thread.join();
            }
        }
    }

private:
    void doAccept() {
        acceptor_.async_accept([this](beast::error_code ec, tcp::socket socket) {
            if (!ec) {
                std::make_shared<AsyncHttpSession>(std::move(socket), router_)->start();
            }
            doAccept(); // 继续接受连接
        });
    }
};


// 示例应用
int start_main() {
    try {
        net::io_context ioc;
        AsyncHttpServer server{ioc, {net::ip::make_address("0.0.0.0"), 9999}};

        // 配置路由
        auto &router = server.getRouter();

        // 快速响应的路由
        router.addRoute("/", HttpMethod::GET, ContentType::ANY, [](const RequestContext &ctx) -> HttpResponse {
            return HttpResponse(200, R"({"message": "Hello World!", "timestamp": ")" +
                                     std::to_string(
                                         std::chrono::duration_cast<std::chrono::milliseconds>(
                                             std::chrono::system_clock::now().time_since_epoch()).
                                         count()) + "\"}");
        });

        // 模拟耗时操作的路由
        router.addRoute("/api/slow", HttpMethod::GET, ContentType::ANY,
                        [](const RequestContext &ctx) -> HttpResponse {
                            // 模拟耗时操作（不会阻塞其他请求）
                            std::this_thread::sleep_for(std::chrono::seconds(3));
                            return HttpResponse(200, R"({"message": "Slow operation completed"})");
                        });

        // 数据处理路由
        router.addRoute("/api/data", HttpMethod::POST, ContentType::JSON,
                        [](const RequestContext &ctx) -> HttpResponse {
                            // 模拟数据处理
                            std::this_thread::sleep_for(std::chrono::milliseconds(500));
                            return HttpResponse(201,
                                                R"({"message": "Data processed", "received": )" + ctx.body + "}");
                        });

        // 查询参数示例
        router.addRoute("/api/users", HttpMethod::GET, ContentType::ANY,
                        [](const RequestContext &ctx) -> HttpResponse {
                            std::string name = ctx.getQueryParam("name", "Anonymous");
                            std::string age = ctx.getQueryParam("age", "0");

                            // 模拟数据库查询
                            std::this_thread::sleep_for(std::chrono::milliseconds(100));

                            return HttpResponse(200,
                                                R"({"users": [{"name": ")" + name + R"(", "age": )" + age + R"(}]})");
                        });

        // 链式路由配置
        router("/api/health", HttpMethod::GET, ContentType::ANY,
               [](const RequestContext &ctx) -> HttpResponse {
                   return HttpResponse(200, R"({"status": "healthy", "uptime": "1000s"})");
               })
        ("/api/version", HttpMethod::GET, ContentType::ANY,
         [](const RequestContext &ctx) -> HttpResponse {
             return HttpResponse(200, R"({"version": "1.0.0", "build": "2024-01-01"})");
         });

        server.run();
    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }

    return 0;
}
